from typing import Sequence

from PIL import Image

from hhutil.io import fmt_path

from paddle.fluid.io import Dataset


class DIV2K(Dataset):

    def __init__(self, lr_root, hr_root, scale, transform=None):
        super().__init__()
        self.lr_root = fmt_path(lr_root)
        self.hr_root = fmt_path(hr_root)
        self.scale = scale
        self.transform = transform

        self.n = len(list(self.hr_root.glob("*.png")))
        self.lr_fps = [
            self.lr_root / ("%04dx%d.png" % (i, scale)) for i in range(1, self.n + 1)
        ]
        self.hr_fps = [
            self.hr_root / ("%04d.png" % i) for i in range(1, self.n + 1)
        ]

    def __getitem__(self, item):
        lr = Image.open(self.lr_fps[item]).convert('RGB')
        hr = Image.open(self.hr_fps[item]).convert('RGB')
        if self.transform is not None:
            lr, hr = self.transform(lr, hr)
        return lr, hr

    def __len__(self):
        return self.n


class CachedDataset(Dataset):

    def __init__(self, dataset, transform=None):
        super().__init__()
        self.dataset = dataset
        self.transform = transform
        self._cache = {}

    def cache(self):
        for i in range(len(self.dataset)):
            self._cache[i] = self.dataset[i]

    def __getitem__(self, item):
        if item in self._cache:
            x = self._cache[item]
        else:
            x = self.dataset[item]
            self._cache[item] = x
        if self.transform is not None:
            if not isinstance(x, Sequence):
                x = [x]
            x = self.transform(*x)
        return x

    def __len__(self):
        return len(self.dataset)